import os
import numpy as np
import joblib
import torch

from .Onehot import Target2text
from .Model import Model

class Adtection:
    def __init__(self) -> None:
        self.model = Model()
        cwd = os.getcwd()
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model1_path = os.path.join(cwd,"Adtection/models/model.pth") 
        model2_path = os.path.join(cwd,"Adtection/models/StandardScaler.pkl")
        if not os.path.exists(model1_path) or not os.path.exists(model2_path):
            model1_path = os.path.join(cwd,"models/model.pth")
            model2_path = os.path.join(cwd,"models/StandardScaler.pkl")       
        model_status = torch.load(model1_path,map_location=device)
        self.model.load_state_dict(model_status)
        self.transer = joblib.load(model2_path)

    def predict(self,data:list) -> float:
        """
            input : list [x1,x2,x3]
            return : str text
        """
        self.time_step = 260
        self.data = np.array(data).ravel()
        if self.data.__len__() < self.time_step:
            return None
        else:
            self.data = self.data[len(self.data)-self.time_step:]
            self.data = self.data.reshape((-1,260))
        self.data = self.transer.transform(self.data).ravel()
        self.data = torch.tensor(self.data,dtype=torch.float32).reshape(-1,1,self.time_step)
        ret = self.model(self.data).argmax(2).item()+1
        text = Target2text(ret)
        return text
if __name__ == "__main__":
    # example
    lstm = Adtection()
    lst = [1 for _ in range(26000)]
    print(lstm.predict(lst))